/*
 * Protein.cpp
 *
 *  Created on: Jul 6, 2013
 *      Author: x
 */

#include <iostream>
#include <iomanip>
#include <fstream>

#include <unistd.h>
#include "Protein.h"
#include "Conformer.h"
#include "oppzip.h"
#include "constant.h"
#include "Microstate.h"
#include <Env.h>

#include <map>
#include <iterator>
#include <cstdlib>
#include <ctime>
#include <cmath>
#include <cstring>
#include <stdexcept>

using namespace std;

namespace mcce {

Protein::Protein():eBase(0.0), eFree(0.0), eTotal(0.0), state(0), pairwise(0), n_free(0), biglist(0), free_res(0), nTotalConf(0), conflist(0) {}

Protein::~Protein() {
	if (state)
		delete[] state;

	if (pairwise) {
		for (size_t i=0; i<nTotalConf; i++) {
			delete[]  pairwise[i];
		}
		delete[] pairwise;
	}

	if (conflist) {
		delete[] conflist;
	}
}

Env Protein::env = Env();

void Protein::loadEnv() {
	env.loadFromRunprm();
}

Protein& Protein::loadConfFromHead3(string fName) {
	ifstream ifile(fName.c_str());
	if (!ifile) {
		cerr << "Open file \"" << fName << "\" fail!" << endl;
		//throw runtime_error("Can't open file " + fName);
		return *this;
	}

	string eachLine;
	getline(ifile, eachLine);


	int nConf = 0;
	while (getline(ifile, eachLine)) {
		Conformer conf;
		conf.loadFromHead3(eachLine);
		nConf++;
		bool resFound = false;
		for (vector<Residue>::iterator ir=residues.begin(); ir != residues.end(); ir++) {
			if (conf.isBelongToRes(*ir)) {
				// conformer belongs to an existing residue.
				ir->addConformer(conf);
				resFound = true;
				break;
			}
		}
		if (!resFound) {
			// Create a new residue.
			Residue newRes(conf);
			residues.push_back(newRes);
		}

	}

	// set the pointer from conformer to residue, and from residue to protein.
	indexAllConf();

	return *this;
}

void Protein::indexAllConf() {
	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		for (vector<Conformer>::iterator confIt=resIt->getConformers().begin(); confIt != resIt->getConformers().end(); confIt++) {
			confIt->setInRes(&(*resIt));
		}
		resIt->setInProt(this);
	}
}

Protein& Protein::adjustHead3Flag() {
	for (vector<Residue>::iterator it=residues.begin(); it != residues.end(); it++) {
		it->adjustHead3Flag();
	}

	// save a copy of all the free residues after the flags of head3 are adjusted.
	getFreeResidues();

	return *this;
}

Protein& Protein::loadConfOpp(string fName, float scale_ele, float scale_vdw) {
	zopp::EMATRIX ematrix;
	ematrix.n = 0;
	zopp::load_energies(&ematrix, ".", 0);

	map<string, int> confIdTable;
	for (int i=0; i<ematrix.n; i++) {
		confIdTable.insert(make_pair(string(ematrix.conf[i].uniqID), i));
	}
	cout << ematrix.n << " conoformers in energies.opp" << endl;
	nTotalConf = ematrix.n;

//	pairwise.clear();
//	pairwise.resize(ematrix.n);
//	for (size_t iz=0; iz<ematrix.n; iz++) {
//		pairwise[iz].resize(ematrix.n);
//	}
	pairwise = new float*[nTotalConf];
	for (size_t i=0; i<nTotalConf; i++) {
		pairwise[i] = new float[nTotalConf];
	}
//	conflist.resize(ematrix.n);
	conflist = new Conformer*[nTotalConf];

	//int checkCounter = 0;
	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		for (vector<Conformer>::iterator confIt=resIt->getConformers().begin(); confIt != resIt->getConformers().end(); confIt++) {
			if (confIdTable.count(confIt->getUniqId())) {
				int confIndex = confIdTable[confIt->getUniqId()];
//				vector<float> confPairwise;
				for (int i=0; i<ematrix.n; i++) {
					// use the same pairwise interaction for every two conformers, so that the pairwise interaction is the same between
					//  conformer i and j and between j and i.
					float sumPairwise = ((ematrix.pw[confIndex][i].ele + ematrix.pw[i][confIndex].ele) * scale_ele
									   + (ematrix.pw[confIndex][i].vdw + ematrix.pw[i][confIndex].vdw) * scale_vdw) / 2;
					Conformer::Pairwise confPw(ematrix.pw[confIndex][i].ele * scale_ele,
											   ematrix.pw[confIndex][i].vdw * scale_vdw,
											   sumPairwise);
					pairwise[confIndex][i] = sumPairwise;

//					confPairwise.push_back(sumPairwise);
					confIt->addConfInteraction(string(ematrix.conf[i].uniqID), confPw);
				}
				conflist[confIndex] = &(*confIt);
//				pairwise[confIndex] = confPairwise;
			}
		}
	}


	zopp::free_ematrix(&ematrix);

	return *this;
}

void Protein::findResNeighbor(float threshold) {
	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		resIt->setNeighborRes(threshold);
		// get all those neighbor residues that are free
		resIt->setFreeNeighborRes();
	}
}



float Protein::calculateBase(float ph, float eh) {
	float monte_temp = 298.15;

	eBase = 0.0;

	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		for (vector<Conformer>::iterator confIt=resIt->getConformers().begin(); confIt != resIt->getConformers().end(); confIt++) {
			confIt->setEPh(monte_temp/ROOMT * confIt->getH() * (ph - confIt->getPKa()) * PH2KCAL);
			confIt->setEEh(monte_temp/ROOMT * confIt->getE() * (eh - confIt->getEm()) * PH2KCAL / 58.0);


			confIt->setESelf0();   // set eSelf0 from all the self energy terms defined in conformer.

			float mfe = 0.0;
			for (vector<Residue>::iterator secondResIt = residues.begin(); secondResIt != residues.end(); secondResIt++) {
				if (!secondResIt->isFixed()) continue;
				for (vector<Conformer>::iterator secondConfIt = secondResIt->getConformers().begin(); secondConfIt != secondResIt->getConformers().end(); secondConfIt++) {
					mfe += confIt->getPairwise()[secondConfIt->getUniqId()].sumPairwise * secondConfIt->getOcc();
				}
			}
			confIt->setEMfe(mfe);
			confIt->setESelf(mfe + confIt->getESelf0());
		}
	}

	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		if (!resIt->isFixed()) continue;
		for (vector<Conformer>::iterator confIt=resIt->getConformers().begin(); confIt != resIt->getConformers().end(); confIt++) {
			eBase += confIt->getESelf() * confIt->getOcc();
		}
	}

	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		if (!resIt->isFixed()) continue;
		for (vector<Conformer>::iterator confIt=resIt->getConformers().begin(); confIt != resIt->getConformers().end(); confIt++) {
			for (vector<Residue>::iterator secondResIt = resIt; secondResIt != residues.end(); secondResIt++) {
				if (!secondResIt->isFixed()) continue;
				for (vector<Conformer>::iterator secondConfIt = secondResIt->getConformers().begin(); secondConfIt != secondResIt->getConformers().end(); secondConfIt++) {
					eBase -= (confIt->getPairwise()[secondConfIt->getUniqId()].sumPairwise
							* confIt->getOcc() * secondConfIt->getOcc());
				}
			}
		}
	}

	return eBase;
}


void Protein::setGroupedFreeConf() {
	for (vector<Residue>::iterator it = residues.begin(); it != residues.end(); it++) {
		it->setGroupedFreeConf();
	}

	findResNeighbor(env.big_pairwise);
}

void Protein::clearConfCounter() {
	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		for (vector<Conformer>::iterator confIt=resIt->getConformers().begin(); confIt != resIt->getConformers().end(); confIt++) {
			confIt->setCounter(0);
		}
	}
}

void Protein::setRandomMcState() {
	vector<Residue*> freeResidues;
	// get all the free residues.
	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		if (!resIt->isFixed()) {
			freeResidues.push_back(&(*resIt));
		}
	}
	mcState.setRandomState(freeResidues);
}

float Protein::calculateEFree() {
	eFree = 0.0;
	int kr, ir;
	int n_free = freeResidues.size();
	/* Triangl calculation, no bias but expensive */
	for (kr=0; kr<n_free; kr++) eFree += conflist[state[kr]]->getESelf();
	for (kr=0; kr<n_free; kr++) {
		for (ir=0; ir<kr; ir++)
			eFree += pairwise[state[kr]][state[ir]];
	}

	return eFree;

}

void Protein::setupMC(float ph, float eh) {
	// get all the free residues, and store them in member freeResidues
	// and get all the free conformers of them, and store them in member freeConformers
	setFreeResConf();


	// set all the neighbors of residues and all the free neighbors.
	findResNeighbor(env.big_pairwise);

	// set the biglist
	map<string, int> resNameToIndex;       // map from residue name to index for free residues
	for (size_t i=0; i<freeResidues.size(); i++) {
		resNameToIndex.insert(make_pair(freeResidues[i]->getResName(), i));
	}

	map<string, int> confNameToIndex;
	for (size_t i=0; i<nTotalConf; i++) {
		confNameToIndex.insert(make_pair(conflist[i]->getUniqId(), i));
	}
//	biglist.clear();
//	for (size_t i=0; i<freeResidues.size(); i++) {
//		vector<int> resN;
//		for (size_t j=0; j<freeResidues[i]->getFreeNeighborRes().size(); j++) {
//			resN.push_back(nameToIndex[freeResidues[i]->getFreeNeighborRes()[j]->getResName()]);
//		}
//		biglist.push_back(resN);
	n_free = freeResidues.size();
	if (free_res) delete[] free_res;

	free_res = new RESIDUE[n_free];
	for (size_t i=0; i<n_free; i++) {
		free_res[i].n = freeResidues[i]->getFreeConformers().size();
		free_res[i].conf = new int[free_res[i].n];
		for (size_t j=0; j<free_res[i].n; j++) {
			free_res[i].conf[j] = confNameToIndex[freeResidues[i]->getFreeConformers()[j]->getUniqId()];
		}
	}

	if (biglist) delete[] biglist;
	biglist = new BIGLIST[n_free];
	for (size_t i=0; i<n_free; i++) {
		biglist[i].n = freeResidues[i]->getFreeNeighborRes().size();
		biglist[i].res = new int[biglist[i].n];
		for (size_t j=0; j<biglist[i].n; j++) {
			biglist[i].res[j] = resNameToIndex[freeResidues[i]->getFreeNeighborRes()[j]->getResName()];
		}
	}

	// calculate the base energy
	calculateBase(ph, eh);
}

void Protein::setState() {
	if (state) delete[] state;
	state = new int[n_free];

	for (size_t i=0; i<n_free; i++) {
		state[i] = free_res[i].conf[rand()/(RAND_MAX/free_res[i].n + 1)];
	}
}

void Protein::mcSample(int nCycles, float ph, float eh){

	nCycles *= freeConformers.size();

	if (n_free == 0) {
		cout << "NO free residue in mcSample" << endl;
		return;
	}

	// reset the counter of all the conformers to 0, but not the occ
	clearConfCounter();


	int mem =freeResidues.size() * sizeof(int);
	int *old_state = new int[freeResidues.size()];


	float bFactor = -KCAL2KT/(env.monte_temp/ROOMT);

	int cycles, n_total, n_cycle;
	if (env.monte_trace > 0) {
		cycles = (nCycles-1)/env.monte_trace + 1; /* round up */
		n_total = cycles*env.monte_trace;
		n_cycle = env.monte_trace;
	}
	else {
		cycles = 1;
		n_total = n_cycle = nCycles;
	}

	float E_minimum, E_state, old_E=0.0, dE=0.0;
	E_minimum = E_state = calculateEFree();

	int old_conf=0, new_conf=0;



	int iflip=0, iconf=0, ires=0, nflips=0;

	time_t timeA = time(NULL);
	for (int i=0; i<cycles; i++) {
//		cout << i * n_cycle << " E_running: " << fixed << setprecision(3) << E_state + getEBase() << endl;
		int repeatTimes = 0;
		bool findNewState = false;
		int iter = n_cycle;
		while(iter) {
			old_E = E_state;
			memcpy(old_state, state, mem);
			ires  = rand()/(RAND_MAX/n_free + 1);
			old_conf = state[ires];
			while (true) {
				iconf = rand()/(RAND_MAX/free_res[ires].n + 1);
				new_conf = free_res[ires].conf[iconf];
				if (old_conf != new_conf) break;
			}


			state[ires] = new_conf;
			E_state += (conflist[new_conf]->getESelf() - conflist[old_conf]->getESelf());
//			E_state += (conflist[new_conf]->eSelf - conflist[old_conf]->eSelf);

			for (size_t j=0; j<n_free; j++) {
				E_state += (pairwise[new_conf][state[j]] - pairwise[old_conf][state[j]]);
			}

			if (rand() & 1) {   /* do multiple flip if odd number */
				if (biglist[ires].n) {
					nflips = env.monte_flips > (biglist[ires].n+1) ?
							(biglist[ires].n+1): env.monte_flips;

					for (int k=1; k<nflips; k++) {
						iflip = biglist[ires].res[rand()/(RAND_MAX/biglist[ires].n + 1)];
						iconf = rand()/(RAND_MAX/free_res[iflip].n + 1);
						old_conf = state[iflip];
						new_conf = free_res[iflip].conf[iconf];

						state[iflip] = new_conf;
						E_state += (conflist[new_conf]->getESelf() - conflist[old_conf]->getESelf());
//						E_state += (conflist[new_conf]->eSelf - conflist[old_conf]->eSelf);

						for (size_t j=0; j<n_free; j++) {
							E_state += (pairwise[new_conf][state[j]] - pairwise[old_conf][state[j]]);
						}
					}
				}
			}



			if (E_minimum > E_state) E_minimum = E_state;
			dE = E_state - old_E;
			if (dE < 0.0) {                                 /* go to new low */
				findNewState = true;
			}
			/*<<< Boltman distribution >>>*/
			else if ((float) rand()/RAND_MAX < exp(bFactor*dE)) { /* Go */
				findNewState = true;
			}
			else {                                                    /* stay, restore the state */
				findNewState = false;
				memcpy(state, old_state, mem);
				E_state = old_E;
				repeatTimes++;
			}

//			for (size_t j=0; j<n_free; j++) {
////				conflist[state[j]]->counter++;
//				conflist[state[j]]->incrementCounter();
//
//			}
			if (findNewState && (repeatTimes != 0)) {
				for (size_t j=0; j<n_free; j++) {
					conflist[old_state[j]]->incrementCounter(repeatTimes);
				}
				repeatTimes = 1;
			}
			iter--;
		}
		if (findNewState) {
			for (size_t j=0; j<n_free; j++) {
				conflist[state[j]]->incrementCounter();
			}
		}
		else {
			for (size_t j=0; j<n_free; j++) {
				conflist[state[j]]->incrementCounter(repeatTimes);
			}
		}
	}
	cout << "exit, fixed: " << getEBase() << " total: " << E_state + getEBase() << endl;

	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		if (resIt->isFixed()) continue;
		for (vector<Conformer>::iterator confIt=resIt->getConformers().begin(); confIt != resIt->getConformers().end(); confIt++) {
			if (confIt->getHFlag() == 't') continue;
			confIt->setOcc((float) confIt->getCounter()/n_total);
		}
	}

	cout << "final random: " << rand() << endl;
	cout << "MC time: " << time(NULL) - timeA << endl;
	delete old_state;

}




void Protein::printConfInfo() {
	ofstream ofile("occ.txt", ios::out);
	ofstream cfile("cm_out", ios::out);
	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		//if (resIt->isFixed()) continue;
		for (vector<Conformer>::iterator confIt=resIt->getConformers().begin(); confIt != resIt->getConformers().end(); confIt++) {
			ofile << confIt->getUniqId();
			for (vector<float>::iterator eachOccIt=confIt->getMocc().begin(); eachOccIt != confIt->getMocc().end(); eachOccIt++) {
				ofile << setiosflags(ios::fixed) << setprecision(3) << ' ' <<  *eachOccIt;
				//if (confIt->getOcc() < 0.01) continue;
				if (confIt->getHFlag() == 't') {
					cfile << setiosflags(ios::fixed) << setprecision(3) << "Fixed " << confIt->getUniqId() << setw(10) << confIt->getESelf0() << setw(10) << confIt->getESelf() << setw(10) << (float) confIt->getOcc() << endl;
				}
				else {
					cfile << setiosflags(ios::fixed) << setprecision(3) << "Free  " << confIt->getUniqId() << setw(10) << confIt->getESelf0() << setw(10) << confIt->getESelf() << setw(10) << (float) confIt->getOcc() << endl;
				}
			}
			ofile << endl;
		}
	}
	ofile.close();
	cfile.close();
}

void Protein::reduceFreeRes() {
	float monte_reduce = env.monte_reduce;
	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		if (resIt->isFixed()) continue;
		for (vector<Conformer>::iterator confIt=resIt->getConformers().begin(); confIt != resIt->getConformers().end(); confIt++) {
			if (confIt->getHFlag() == 'f' && confIt->getOcc() < monte_reduce) {
				confIt->setHFlag('t');
				confIt->setOcc(0.0);
			}
		}
		// readjust the flag and occ of conformers in this residue
		resIt->adjustHead3Flag();
	}
}

int Protein::getNumberOfFreeConf() {
	int counter = 0;

	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		if (resIt->isFixed()) continue;
		for (vector<Conformer>::iterator confIt=resIt->getConformers().begin(); confIt != resIt->getConformers().end(); confIt++) {
			if (confIt->getHFlag() == 't') continue;
			counter++;
		}
	}

	return counter;
}

vector<Microstate>& addOneResToMState(vector<Microstate>&mStates, vector<Conformer*> &confs){
	if (mStates.size() == 0) {
		for (vector<Conformer*>::iterator eachConf=confs.begin(); eachConf != confs.end(); eachConf++) {
			vector<Conformer*> newConfs;
			newConfs.push_back(*eachConf);
			Microstate newState(newConfs);
			mStates.push_back(newState);
		}
	}
	else {
		vector<Microstate> bkStates(mStates);
		mStates.clear();
		for (vector<Microstate>::iterator eachState=bkStates.begin(); eachState != bkStates.end(); eachState++) {
			for (vector<Conformer*>::iterator eachConf=confs.begin(); eachConf != confs.end(); eachConf++) {
				vector<Conformer*> newConfs(eachState->getConformers());
				newConfs.push_back(*eachConf);
				Microstate newState(newConfs);
				mStates.push_back(newState);
			}
		}
	}

	return mStates;
}

void Protein::enumMstate() {
	vector<Microstate> allStates;
	vector<Residue*> freeResidues(getFreeResidues());

	size_t n_free = freeResidues.size();

	if (n_free ==0) return;

	Microstate newState;
	for (vector<Residue*>::iterator resIt=freeResidues.begin(); resIt != freeResidues.end(); resIt++) {
		Conformer *firstConf = &((*resIt)->getConformers()[0]);
		newState.addConf(firstConf);
	}

	float E_state = newState.calculateEState();
	float E_min = E_state;

	allStates.push_back(newState);

	size_t ires = 0;
	Conformer *oldConf, *newConf;

	while (1) {
		ires = 0;
		oldConf = newState.confAt(ires);
		freeResidues[ires]->setOnFreeConf(freeResidues[ires]->getOnFreeConf() + 1);

		while (freeResidues[ires]->getOnFreeConf() >= freeResidues[ires]->getFreeConformers().size()) {
			freeResidues[ires]->setOnFreeConf(0);
			newConf =  &(freeResidues[ires]->getConformers()[0]);

			newState.setConfAt(ires, newConf);

			E_state += newConf->getESelf() - oldConf->getESelf();
			for (vector<Residue*>::size_type jres=0; jres<freeResidues.size(); jres++) {
				E_state += (newConf->getPairwise()[newState.confAt(jres)->getUniqId()].sumPairwise
						  - oldConf->getPairwise()[newState.confAt(jres)->getUniqId()].sumPairwise);
			}

			ires++;

			// Have iterated the final residue.
			if (ires == n_free) break;

			oldConf = newState.confAt(ires);

			freeResidues[ires]->setOnFreeConf(freeResidues[ires]->getOnFreeConf() + 1);
		}

		if (ires == n_free) break;

		newConf = &(freeResidues[ires]->getConformers()[freeResidues[ires]->getOnFreeConf()]);

		newState.setConfAt(ires, newConf);

		E_state += newConf->getESelf() - oldConf->getESelf();
		for (vector<Residue*>::size_type jres=0; jres<freeResidues.size(); jres++) {
			E_state += (newConf->getPairwise()[newState.getConformers()[jres]->getUniqId()].sumPairwise
					  - oldConf->getPairwise()[newState.getConformers()[jres]->getUniqId()].sumPairwise);
		}

		newState.setEstate(E_state);
		allStates.push_back(newState);
		if (E_state < E_min) E_min = E_state;
	}
	cout << "Number of microstates: " << allStates.size() << endl;
	//float bFactor = -KCAL2KT/(env.monte_temp/ROOMT);
	float bFactor = -KCAL2KT/(env.monte_temp/ROOMT);

	for (vector<Residue*>::iterator eachResIt=freeResidues.begin(); eachResIt!=freeResidues.end(); eachResIt++) {
		for (vector <Conformer>::iterator eachConfIt=(*eachResIt)->getConformers().begin(); eachConfIt!=(*eachResIt)->getConformers().end(); eachConfIt++) {
			if (eachConfIt->getHFlag() == 't') continue;
			eachConfIt->setOcc(0.0);
		}
	}

	float totOcc = 0.0;
	for (vector<Microstate>::iterator eachStateIt=allStates.begin(); eachStateIt != allStates.end(); eachStateIt++) {
		totOcc += exp(bFactor * (eachStateIt->getEState() - E_min));
	}

	ofstream tempfp("temp.txt", ios::out);
	for (vector<Microstate>::iterator eachStateIt=allStates.begin(); eachStateIt != allStates.end(); eachStateIt++) {
		float partialOcc = exp(bFactor * (eachStateIt->getEState() - E_min))/totOcc;
		if (partialOcc > 0.001) tempfp << "occ: " << partialOcc << endl;
		for (vector<Conformer*>::iterator eachConfIt=eachStateIt->getConformers().begin(); eachConfIt!=eachStateIt->getConformers().end(); eachConfIt++) {
			(*eachConfIt)->setOcc((*eachConfIt)->getOcc() + partialOcc);
		}
	}
	cout << "final" << endl;
}

void Protein::saveConfOcc() {
	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		//if (resIt->isFixed()) continue;
		for (vector<Conformer>::iterator confIt=resIt->getConformers().begin(); confIt != resIt->getConformers().end(); confIt++) {
			confIt->getMocc().push_back(confIt->getOcc());
		}
	}
}

void Protein::setFreeResConf() {
	freeResidues.clear();
	freeConformers.clear();
	for (vector<Residue>::iterator resIt=residues.begin(); resIt != residues.end(); resIt++) {
		if (!resIt->isFixed()) {
			freeResidues.push_back(&(*resIt));
			resIt->setFreeConformers();
			freeConformers.insert(freeConformers.end(), resIt->getFreeConformers().begin(), resIt->getFreeConformers().end());
		}
	}
}
} /* namespace mcce */
